package com.yanggu.flink.datastream_api.state.keyedstate

import com.yanggu.flink.datastream_api.pojo.Event
import com.yanggu.flink.datastream_api.source.ClickSource
import org.apache.flink.api.common.functions.AggregateFunction
import org.apache.flink.api.common.state.{AggregatingState, AggregatingStateDescriptor, ValueState, ValueStateDescriptor}
import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.KeyedProcessFunction
import org.apache.flink.streaming.api.scala._
import org.apache.flink.util.Collector

/**
 * 对用户点击事件流每 5 个数据统计一次平均时间戳。
 * 这是一个类似计数窗口（CountWindow）求平均值的计算
 */
object AggregatingStateDemo {

  def main(args: Array[String]): Unit = {
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    env.setParallelism(1)

    env
      .addSource(new ClickSource)
      .keyBy(data => data.getName)
      .process(new AvgTsResult)
      .print()

    env.execute()

  }

}

class AvgTsResult extends KeyedProcessFunction[String, Event, String] {

  private var aggState: AggregatingState[Event, Long] = _

  private var countState: ValueState[Int] = _

  override def open(parameters: Configuration): Unit = {
    val aggStateDescriptor = new AggregatingStateDescriptor[Event, (Long, Long), Long]("agg", new AggregateFunction[Event, (Long, Long), Long] {

      override def createAccumulator() = (0L, 0L)

      override def add(value: Event, accumulator: (Long, Long)) = (accumulator._1 + value.getTimestamp, accumulator._2 + 1)

      override def getResult(accumulator: (Long, Long)) = accumulator._1 / accumulator._2

      override def merge(a: (Long, Long), b: (Long, Long)) = (a._1 + b._1, a._2 + b._2)

    }, classOf[(Long, Long)])

    aggState = getRuntimeContext.getAggregatingState(aggStateDescriptor)

    val countDescriptor = new ValueStateDescriptor[Int]("count", classOf[Int])

    countState = getRuntimeContext.getState(countDescriptor)

  }

  override def processElement(value: Event, ctx: KeyedProcessFunction[String, Event, String]#Context, out: Collector[String]): Unit = {
    //记录次数
    val count = countState.value() + 1
    countState.update(count)

    //添加aggState中
    aggState.add(value)

    if (count == 5) {
      out.collect(s"当前用户: ${ctx.getCurrentKey}, 每隔5次访问时间戳的平均值: ${aggState.get()}")
      aggState.clear()
      countState.clear()
    }
  }

}
